package org.example.rabbitmq.consumer.rate;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;

/**
 * 队列消费限流器
 *
 * 可以用于单进程、一个队列与Consumer不一一对应的情况，比如：一个队列由多个Consumer消费
 */
public class RamRateLimiter implements RateLimiter {
    Map<String, Semaphore> semaphoreMap = new ConcurrentHashMap<>();

    public RamRateLimiter(String... queueNames) {
        for (String queueName : queueNames) {
            addQueueRateLimiter(queueName);
        }
    }

    @Override
    public void acquire(String queueName) throws InterruptedException {
        semaphoreMap.get(queueName).acquire();
    }

    @Override
    public void release(String queueName) {
        semaphoreMap.get(queueName).release();
    }

    @Override
    public void addQueueRateLimiter(String queueName) {
        semaphoreMap.put(queueName, new Semaphore(1));
    }
}
